"""
该模块有两个函数
函数一、
用于处理sdf数据,对Z做二维傅里叶变换，滤出在center邻域内，allow_error范围内的波.
提供函数
fft_analyse_filter(Z,X,Y=0,
    data_dim=data_dim,one2two=0,allow_error=0.1,ifpositive=0,center=lp.wavenumber)
注意：
    1.Z是数据数组，且Z是经过转置之后的矩阵(与X，Y是直接由meshgrid函数的结果相对应)
    2.采样点间距依赖"读取数据"块，值为dx = X[0, 1] - X[0, 0];dy = Y[1, 0] - Y[0, 0]

需要定义的参数：
    1.ifpositive
    ·等于0时，直接np.fft,fft频率为
    f = [0, 1, ...,   n/2-1,     -n/2, ..., -1] / (d*n)(n为偶数)
    f = [0, 1, ..., (n-1)/2, -(n-1)/2, ..., -1] / (d*n)(n是奇数)的形式,共n项
    这里的输出经过np.fft.fftshift的变换
    f = [     -n/2, ..., -1,0, 1, ...,   n/2-1,] / (d*n)
    并且滤波时，频率先去绝对值处理(这样的，傅里叶反演之后自然是实数值)
    ·等于1时，采用频率为f = [0, 1, ...,   n/2-1,     n/2, ..., n-1] / (d*n)
    ·如上所示，频率的实际含义为波数
    2.allow_error:频率差大于allow_error*center频率的部分将被删去
    
    3.center 默认为al.sp.wavenumber
    
    4.one2two

返回值为 二维：(var_ifft2, var_fft2, var_x_fre, var_y_fre)
返回值为 一维：(var_ifft, X, var_fft, var_x_fre)
#=======<<<函数一描述结束=============================================================================
函数二、

"""
#=================test 以下为测试脚本所用的自定义数据==================
# data_dim = 2
# one2two = 1
# X = np.linspace(0, 1e-5, 500)
# Y = np.linspace(-1e-5, 1e-5, 500)
# X, Y = np.meshgrid(X, Y)
# [rows, cols] = X.shape
# print("[rows, cols] =", rows, cols)
# Z = np.zeros([rows, cols])
# for i in range(rows):
#     for j in range(cols):
#         # Z[i, j] = al.mathfunc.spherical_wave(0.4e-12,X[i,j],Y[i,j],0 )
#         Z[i, j] = al.mathfunc.spherical_wave(0, X[i, j], Y[i, j], 0)
# # ;;;;;
# data_dim = 1
# one2two = 0
# X = np.linspace(-1e-5, 1e-5, 1000)
# Z = np.zeros(len(X))
# for i in range(len(X)):
#     Z[i] = al.mathfunc.plane_wave(0, X[i], 0, 0)
#精确分解测试
#2d=========
# data_dim = 2
# one2two = 1
# X = np.linspace(0, 4 * np.pi, 164, endpoint=0)
# Y = np.linspace(0, 4 * np.pi, 164, endpoint=0)  #Y的范围需要自己调整
# X, Y = np.meshgrid(X, Y)
# Z = np.cos(X + Y)
# al.lp.wavenumber = np.sqrt(2) / 2 / np.pi
# al.pp.vmax = 1
# al.pp.vmin = -1
#1d==========
# data_dim = 1
# one2two = 1
# X = np.linspace(0, 4 * np.pi, 164, endpoint=0)
# Y = np.linspace(0, 4 * np.pi, 164, endpoint=0)  #Y的范围需要自己调整
# Z = np.cos(X)
# al.lp.wavenumber = 1 / 2 / np.pi
# al.pp.vmax = 1
# al.pp.vmin = -1
import matplotlib as mpl
# mpl.use('Agg')
# import sdf_helper as sh
import matplotlib.pyplot as plt
import numpy as np
from numpy.core.numeric import Inf
from scipy import fft as scpfft
from scipy import signal as scpsgnl
from ..mathfunc import laser_parm as lp
from ..tool import code_help
from ..sim_parm import *


def fft_analyse_filter(Z,
                       X,
                       Y=0,
                       data_dim=None,
                       one2two=0,
                       allow_error=Inf,
                       ifpositive=0,
                       center=None):
    if (center == None):
        if (allow_error == Inf):
            print('using allow_error=inf,center=1')
            center = 1
        else:
            if (allow_error == Inf):
                print("not allowed 'allow_error', please specify it")
                return 0
    if data_dim == 2:
        print("choosing data_dim == 2")
        # ========================1d=======================================
        return data_dim_2(Z,
                          X,
                          Y,
                          allow_error=allow_error,
                          ifpositive=ifpositive,
                          center=center)

    elif (data_dim == 1):
        print("choosing data_dim == 1")
        return data_dim_1(Z,
                          X,
                          Y=Y,
                          one2two=one2two,
                          allow_error=allow_error,
                          ifpositive=ifpositive,
                          center=center)
    else:
        print('请指定正确的data_dim')
        return 0


#=======================================================================
def data_dim_2(Z, X, Y, allow_error, ifpositive, center):
    #===============参数设置=================================
    # dx = (xmax - xmin) / nx
    # dy = (ymax - ymin) / ny
    dx = X[0, 1] - X[0, 0]
    dy = Y[1, 0] - Y[0, 0]
    #===============画出原图==================================
    # ax = plt.subplot(1, 2, 1)
    # tempax = al.plot.plot_2d(Z, X, Y, ax=ax, ifsdf=0)
    #===============开始滤波，并画图==================================
    #========================2d=======================================
    var_fft2 = np.fft.fft2(Z)
    code_help.dprint("var_fft2", var_fft2)
    code_help.dprint('len of var_fft2', len(var_fft2[1]), len(var_fft2[:, 1]))
    #注释np.fft.fftfreq返回值f形式为，注意和实际的k之间相差2pi
    #f = [0, 1, ...,   n/2-1,     -n/2, ..., -1] / (d*n)   if n is even 共n项
    #f = [0, 1, ..., (n-1)/2, -(n-1)/2, ..., -1] / (d*n)   if n is odd  共n项
    if ifpositive == 0:
        print("choosing ifpositive=0")
        var_x_fre = np.fft.fftfreq(len(Z[0]), dx)
        # code_help.dprint("var_x_fre", var_x_fre)
        var_y_fre = np.fft.fftfreq(len(Z[:, 0]), dy)
        # code_help.dprint("var_y_fre", var_y_fre)
        # var_x_fre = np.fft.fftfreq(smp.y - smp.x, X[0][1] - X[0][0])
        # var_y_fre = np.fft.fftfreq(smp.yy - smp.yx, Y[1][0] - Y[0][0])
    elif ifpositive == 1:
        for i in range(len(var_x_fre)):
            var_x_fre[i] = i / (len(Z[0]) * dx)
            var_y_fre[i] = i / (len(Z[:, 0]) * dy)
    # code_help.dprint('in given data,x-max frequency=%f,y-max frequency=%f',
    #      %(var_x_fre[np.argmax(var_fft)],var_y_fre[np.argmax(var_fft)]))
    # code_help.dprint("test==", var.grid_mid.data[1][1] - var.grid_mid.data[1][0],
    #       Y[1][0] - Y[0][0])
    for i in range(len(Z[0])):
        for j in range(len(Z[:, 0])):
            if (np.absolute(
                    np.sqrt(var_x_fre[i]**2 + var_y_fre[j]**2) -
                    center)) > (allow_error * (center)):
                var_fft2[j, i] = 0
    var_ifft2 = np.real(np.fft.ifft2(var_fft2))
    return (var_ifft2, np.fft.fftshift(var_fft2), np.fft.fftshift(var_x_fre),
            np.fft.fftshift(var_y_fre))


#&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&
def data_dim_1(Z, X, Y, one2two, allow_error, ifpositive, center):
    """
        基本和二维一样，不同点有：
            1.
    """
    #===============扩展到二维,画出原图==================================
    if one2two == 1:
        #关键参数的定义
        dx = X[1] - X[0]
        #================================
        initZ = Z
        var_fft = np.fft.fft(initZ)
        # print("var_fft", var_fft)
        ifpositive = 0
        code_help.dprint('len of var_fft', len(var_fft))
        #注释np.fft.fftfreq返回值f形式为，注意和实际的k之间相差2pi
        #f = [0, 1, ...,   n/2-1,     -n/2, ..., -1] / (d*n)   if n is even 共n项
        #f = [0, 1, ..., (n-1)/2, -(n-1)/2, ..., -1] / (d*n)   if n is odd  共n项
        if ifpositive == 0:
            print("choosing ifpositive=0")
            var_x_fre = np.fft.fftfreq(len(initZ), dx)  #1167136
        elif ifpositive == 1:
            print("choosing ifpositive=1")
            var_x_fre = np.zeros(len(var_fft))
            for i in range(len(var_fft)):
                var_x_fre[i] = i / (len(var_fft) * dx)

        print('in given data, max frequency(波数): ',
              var_x_fre[np.argmax(var_fft)], '(乘2pi为：)',
              2 * np.pi * var_x_fre[np.argmax(var_fft)])

        # plt.scatter(var_x_fre, var_fft)
        # plt.show()
        # code_help.dprint("test==", var.grid_mid.data[1][1] - var.grid_mid.data[1][0],
        #       Y[1][0] - Y[0][0])
        if ifpositive == 0:
            for i in range(len(initZ)):
                if (np.absolute(abs(var_x_fre[i]) - center)) > np.absolute(
                        allow_error * (center)):
                    var_fft[i] = 0
        elif ifpositive == 1:
            for i in range(len(initZ)):
                if (np.absolute(abs(var_x_fre[i]) - center)) > np.absolute(
                        allow_error * (center)):
                    var_fft[i] = 0
        # var_ifft = np.real(np.fft.ifft(var_fft))
        var_ifft = np.fft.ifft(var_fft)
        #画图,先把var_ifft转化到二维
        [var_ifft2, X, Y] = code_help.help_one2two(var_ifft, X, Y)
        return (var_ifft2, X, Y, var_fft, var_x_fre)
    elif one2two == 0:
        #关键参数的定义
        # code_help.dprint(X)
        dx = X[1] - X[0]
        #================================
        initZ = Z
        var_fft = np.fft.fft(initZ)
        #加窗
        #w = scpsgnl.blackman(len(initZ))
        #var_fft = scpfft.fft(initZ * w)
        if ifpositive == 0:
            print("choosing ifpositive=0")
            var_x_fre = np.fft.fftfreq(len(initZ), dx)  #1167136
        elif ifpositive == 1:
            print("choosing ifpositive=1")
            var_x_fre = np.zeros(len(var_fft))
            for i in range(len(var_fft)):
                var_x_fre[i] = i / (len(var_fft) * dx)
        print('in given data, max frequency(波数): ',
              var_x_fre[np.argmax(var_fft)], '(乘2pi为：)',
              2 * np.pi * var_x_fre[np.argmax(var_fft)])
        if ifpositive == 0:
            for i in range(len(initZ)):
                if (np.absolute(abs(var_x_fre[i]) - center)) > np.absolute(
                        allow_error * (center)):
                    var_fft[i] = 0
        elif ifpositive == 1:
            for i in range(len(initZ)):
                if (np.absolute(abs(var_x_fre[i]) - center)) > np.absolute(
                        allow_error * (center)):
                    var_fft[i] = 0
        # var_ifft = np.real(np.fft.ifft(var_fft))
        var_ifft = np.fft.ifft(var_fft)
        return (var_ifft, X, np.fft.fftshift(var_fft),
                np.fft.fftshift(var_x_fre))


def ft_continue(
    Z,
    X,
    Y=0,
    data_dim=None,
    kx=None,
    ky=None,
    kz=None,
):
    '''
    数学原理上，是要对Z做无穷区间的连续傅里叶变换，所以要求Z在所给区间外值几乎为0
    傅里叶变化系数取 根号2Pi
    '''
    if (data_dim == 2):
        if ((kx == None) | (ky == None)):
            print('请指定kx和ky')
            return 0
        else:
            code_help.dprint('注意傅里叶系数取 根号2Pi')
            res = 0
            dx = X[0, 1] - X[0, 0]
            dy = Y[1, 0] - Y[0, 0]
            for i in range(Z.shape[0]):
                for j in range(Z.shape[1]):
                    res = res + np.exp(-1j * kx * X[i, j] -
                                       1j * ky * Y[i, j]) * Z[i, j] * dx * dy
            return res / np.sqrt(2 * np.pi)
    elif (data_dim == 1):
        res = 0
        dx = X[1] - X[0]
        code_help.dprint('间距dx为:', dx, '     kx为：', kx)
        code_help.dprint('X为：', X, '    Z为：', Z)
        for i in range(Z.shape[0]):
            res = res + np.exp(-1j * kx * X[i]) * Z[i] * dx
        return res / np.sqrt(2 * np.pi)
    else:
        print('请输入正确的维数--没有三维')
        return 0
    return 0